from kuai.mol import *
from kuai.molio import *
from kuai.element import *

class DffMsdFileReader(MoleculeReader):
    def read(self, file):
        pbc = None
        for i in file:
            try:
                if i.startswith("PBC:"):
                    tokens = i[4:].split()
                    a, b, c = float(tokens[0]), float(tokens[1]), float(tokens[2])
                    alpha, beta, gamma = float(tokens[3]), float(tokens[4]), float(tokens[5])
                    pbc = PBC(a, b, c, alpha, beta, gamma)
                    continue
                else:
                    natoms = int(i.strip())
                    break
            except ValueError:
                pass
            
        atoms = []
        for i in range(natoms):
            i = file.next()
            tokens = i.split()
            if len(tokens) > 8:
                symbol = get_symbol(int(tokens[2].strip()))
                x, y, z = float(tokens[6]), float(tokens[7]), float(tokens[8]) 
                a = Atom(symbol, XYZ(x, y, z))
                a.partial_charge = float(tokens[5])
                a.type = tokens[3]
                if tokens[4] != '?':
                    a.dfftype = tokens[4]
                    
                if len(tokens) >= 12:
                    a.residure = tokens[10]
                    gid = int(tokens[11])
                    if gid != 0:
                        a.group_id = gid
                atoms.append(a)
            else:
                raise IOError('Error Atom Type at %s' % i)
            
        nbonds = int(file.next().strip())
        bonds = []
        for i in range(nbonds):
            tokens = file.next().split()
            i1 = int(tokens[0])
            i2 = int(tokens[1])
            order = int(tokens[2])
            if order == 1:
                order = SINGLE_BOND
            elif order == 2:
                order = DOUBLE_BOND
            elif order == 3:
                order = TRIPLE_BOND
            elif order == -2:
                order = PARTIAL_BOND
            else:
                raise ValueError("Invalid Bond Order " + str(order))
            b = Bond(atoms[i1-1], atoms[i2-1], order)
            bonds.append(b)
        
        try:
            i = file.next()
            if i.startswith('#Formal Charge'):
                n = int(file.next().strip())
                for i in range(n):
                    tokens = file.next().strip().split()
                    a, e = int(tokens[0]), int(tokens[1])
                    atoms[a-1].formal_charge = e
        except StopIteration:
            pass
        result = Molecule(atoms, bonds)
        if pbc:
            result.pbc = pbc
        return result
    
reader = DffMsdFileReader()

class DffMsdFileWriter(MoleculeWriter):
    def write(self, file, v):
        file.write("""#DFF:MSD
#Model Structure Data File    Energy = 0.0
""")
        charged_atoms = []
        file.write(str(len(v.atoms)) + "\n")
        for i in v.atoms:
            id = v.index(i)
            file.write("%6d %8s %3d %12s %7s %6.3f %10.4f %10.4f %10.4f 1 UNK 0\n" \
                       % (id+1, i.symbol() + str(id), i.number(), i.type, "?", \
                          i.partial_charge, i.coords.x, i.coords.y, i.coords.z))
            if i.formal_charge != 0:
                charged_atoms.append(i)
        
        file.write(str(len(v.bonds)) + "\n")
        for i in v.bonds:
            id1 = v.index(i.atom1)+1
            id2 = v.index(i.atom2)+1
            if i.order == SINGLE_BOND:
                order = 1
            elif i.order == DOUBLE_BOND:
                order = 2
            elif i.order == TRIPLE_BOND:
                order = 3
            elif i.order == PARTIAL_BOND:
                order = -2
            else:
                raise ValueError("Invalid Bond Order " + str(order))
            file.write("%6d %6d %3d\n" % (id1, id2, order))
            
        if len(charged_atoms) > 0:
            sum_charge = sum([i.formal_charge for i in charged_atoms], 0)
            file.write("#Formal Charge = " + str(sum_charge) + "\n")
            file.write(str(len(charged_atoms))+"\n")
            for i in charged_atoms:
                file.write("%6d %6d\n" % (v.index(i)+1, i.formal_charge))
        file.write("M  END\n")
        return True

writer = DffMsdFileWriter()
